local bit = require("bit")

local bor = bit.bor
local band = bit.band
local lshift = bit.lshift
local rshift = bit.rshift
local abs = math.abs

local function put2(state, b)
	local word = bor(lshift(state.word, 1), b)
	local left = state.left - 1
	
	state.word = word
	state.left = left

	if left == 0 then
		local next = state.n + 1
		state.buffer[next] = word
		state.n = next
		state.word = 0
		state.left = 32
	end
end

local function putint(state, i, bits)
	local left = state.left
	local word = state.word

	if left > bits then
		-- bits must be < 32 here since left is <= 32
		state.word = bor(lshift(word, bits), i)
		state.left = left - bits
	else
		-- left is >= 1, therefore leftover is <= 31.
		-- left can be 32, but then word is 0, so it doesn't matter how much
		-- it's shifted.
		local leftover = bits - left
		local wword = bor(lshift(word, left), rshift(i, leftover))

		local next = state.n + 1
		state.buffer[next] = wword
		state.n = next

		local newleft = 32 - leftover  -- 32 - [0, 31] === [1, 32]
		state.word = i
		state.left = newleft
	end
end

local function getint(state, i, bits)
	local left = state.left
	local word = state.word

	if left > bits then
		state.word = lshift(word, bits)
		state.left = left - bits
		return rshift(word, 32-bits)
	else
		local leftover = bits - left

		local next = state.n + 1
		local newword = state.buffer[next]
		state.n = next

		state.word = lshift(newword, leftover)
		return bor(rshift(word, 32 - bits), rshift(newword, 32 - leftover))
	end
end

local function getint2(state, i, bits)
	local left = state.left -- left in [0, 31]
	local word = state.word

	-- bits assumed in [1, 32]

	if left >= bits then
		-- bits in [1, 31] here
		state.word = lshift(word, bits)
		state.left = left - bits
		return rshift(word, 32-bits)
	else
		local bitstogo = bits - left -- bitstogo in [1, 32]

		local next = state.n + 1
		local newword = state.buffer[next]
		state.n = next

		state.word = lshift(lshift(newword, bitstogo-1), 1) -- shift by [1, 32]
		return bor(rshift(word, 32 - bits), rshift(newword, 32 - bitstogo))
	end
end

local function Bitstream()
	local self = {word = 0, left = 32, buffer = {}, n = 0}
	setmetatable(self, {__index = {put = put2, putint = putint}})
	return self
end

do
	local bstate = {word = 0, left = 32, buffer = {}, n = 0}

	local start = os.clock()
	for i=1,100000000 do
		put2(bstate, bit.band(i, 1))
	end
	local stop = os.clock()

	print(stop - start, #bstate.buffer)
end

do
	local bitstream = Bitstream()

	local start = os.clock()
	for i=1,100000000 do
		bitstream:putint(bit.band(i, 1), 1)
	end
	local stop = os.clock()

	print(stop - start, #bitstream.buffer)
end